#!/usr/bin/python
import sys

class ObjFile(object):
    """a lnl obj file
    """
    SEP = " "
    MAGICNUMBER = "LINK"
    
    def __init__(self):
        pass
    
    def read(self, f):
        """read from a f
        f should be a iteratable which yield lines.
        """
        wrapper = InputWrapper(f)
        #magic number line
        try:
            self.readmagicnumber(wrapper)
        except StopIteration,si:
            raise Exception("eof while read magic number line.")
        #item number line
        try:
            self.readitemnumberline(wrapper)
        except StopIteration,si:
            raise Exception("eof while read item number line.")
        #segments
        try:
            self.readsegments(wrapper)
        except StopIteration,si:
            raise Exception("eof while read segments.")
        #symbols
        try:
            self.readsymbols(wrapper)
        except StopIteration,si:
            raise Exception("eof while read symbols.")
        #rels
        try:
            self.readrels(wrapper)
        except StopIteration,si:
            raise Exception("eof while read relocations.")
        #data
        try:
            self.readdata(wrapper)
        except StopIteration,si:
            raise Exception("eof while read data.")
    
    def write(self, f):
        #magic number line
        self.writemagicnumber(f)
        #item number line
        self.writeitemnumberline(f)
        #segments
        self.writesegments(f)
        #symbols
        self.writesymbols(f)
        #rels
        self.writerels(f)
        #data
        self.writedata(f)
    
    def readmagicnumber(self, f):
        self.magicnumber = f.next()
        if self.magicnumber != ObjFile.MAGICNUMBER:
            raise Exception("is not a lnl obj file, magic number invalid.")

    def writemagicnumber(self, f):
        print >> f, self.magicnumber
    
    def readitemnumberline(self, f):
        line = f.next()
        line = line.split(ObjFile.SEP)
        if len(line) < 3:
            raise Exception("(%d): item number line has fewer than 3 fields." % f.lineno)
        self.nsegs, self.nsyms, self.nrels = map(int, line[:3])

    def writeitemnumberline(self, f):
        print >> f, ObjFile.SEP.join(map(str, [self.nsegs, self.nsyms, self.nrels]))
    
    def readsegments(self, f):
        self.segs = []
        for i in xrange(1, self.nsegs+1):
            line = f.next()
            line = line.split(ObjFile.SEP)
            if len(line) != 4:
                raise Exception("(%d): segment line must have 4 fileds." % f.lineno)
            name , saddr, length, desc = line
            saddr = int(saddr, 16)
            length = int(length, 16)
            self.segs.append(Segment(i, name, saddr, length, desc))

    def writesegments(self, f):
        for seg in self.segs:
            print >> f, ObjFile.SEP.join(map(str,
                [seg.name, seg.saddr, seg.len,
                map(lambda x:x[1] if x[0] else "", zip([seg.readable, seg.writable, seg.present],"RWP"))]
                ))
    
    def readsymbols(self, f):
        self.syms = []
        for i in xrange(1, self.nsyms+1):
            line = f.next()
            line = line.split(ObjFile.SEP)
            if len(line) != 4:
                raise Exception("(%d): symbol line must have 4 fileds." % f.lineno)
            name , value, seg, stype = line
            value = int(value, 16)
            seg = int(seg, 16)
            self.syms.append(Symbol(i, name, value, seg, stype))
            
    def writesymbols(self, f):
        for sym in self.syms:
            print >> f, ObjFile.SEP.join(map(str,
                [sym.name, sym.value, sym.seg, sym.type]
                ))
    
    def readrels(self, f):
        self.rels = []
        for i in xrange(1, self.nrels+1):
            line = f.next()
            line = line.split(ObjFile.SEP, 4)
            if len(line) < 4:
                raise Exception("(%d): relocation line must have at least 4 fileds." % f.lineno)
            loc, seg, ref, rtype = line[:4]
            if len(line) > 4:
                extra = line[4]
            else:
                extra = None
            seg = int(seg, 16)
            ref = int(ref, 16)
            self.rels.append(Reloc(i, loc, seg, ref, rtype, extra))

    def writerels(self, f):
        for rel in self.rels:
            print >> f, ObjFile.SEP.join(map(str,
                [rel.loc, rel.seg, rel.ref, rel.type, rel.extra if rel.extra else "" ]
                ))
    
    def readdata(self, f):
        try:
            for seg in (x for x in self.segs if x.present):
                line = f.next()
                if len(line) != seg.len * 2:
                    raise Exception("(%d): invalid length in segment %d." % (f.lineno, seg.no))
                seg.data = line
        except StopIteration,si:
            raise Exception("eof while read data.")
    
    def writedata(self, f):
        for seg in (x for x in self.segs if x.present):
            print >> f, seg.data

class InputWrapper(object):
    """wrap a iteraterable,
    to strip ending \n and to record line number
    """
    def __init__(self, f):
        self.f = f
        self.lineno = 0
    
    def __iter__(self):
        return self
    
    def next(self):
        line = self.f.next().strip("\n")
        self.lineno += 1
        return line

class Segment(object):
    """a segment in obj file
    """
    def __init__(self, no, name, saddr, length, desc):
        self.no = no
        self.name = name
        self.saddr = saddr
        self.len = length
        self.readable = "R" in desc 
        self.writable = "W" in desc 
        self.present = "P" in desc 

class Symbol(object):
    """a symbol in obj file
    """
    def __init__(self, no, name, value, seg, stype):
        self.no = no
        self.name = name
        self.value = value
        self.seg = seg
        self.type = stype

class Reloc(object):
    """a relocation in obj file
    """
    def __init__(self, no, loc, seg, ref, rtype, extra):
        self.no = no
        self.loc = loc
        self.seg = seg
        self.ref = ref
        self.type = rtype
        self.extra = extra

def main():
    filedata = \
    """LINK
3 2 1
.text 1000 2500 RP
.data 4000 C00 RWP
.bss 5000 1900 RW
main 80004800 1 D
snprintf 98888888 1 U
09 2 3 A4
""" + \
    "F8" * 0x2500 + "\n" + \
    "07" * 0xC00 + "\n"
    import StringIO
    f=StringIO.StringIO(filedata)
    obj = ObjFile()
    obj.read(f)
    fw=StringIO.StringIO()
    obj.write(fw)
    assert(filedata == f.buf)
    return 0

if __name__ == '__main__':
    sys.exit(main())

